#include "graph_utils.h"
#include <gflags/gflags.h>
#include <butil/logging.h>
#include <brpc/server.h>
#include <brpc/channel.h>
#include <json2pb/pb_to_json.h>
#include <algorithm>
#include <cstring>
#include <functional>
#include <string>
#include <vector>

#include <thread>
#include <condition_variable>

#include <butil/logging.h>
#include "common_unit.h"

namespace rap {

class DocaDmaServiceImpl final : public DocaDmaService { //CPU侧进入这里
public:
    DocaDmaServiceImpl() {}
    ~DocaDmaServiceImpl() {
        if (this->if_pipeline) {
            if (this->backgroundThread.joinable()) {
                backgroundThread.join();
            }
        }
        else{
        #if defined(__aarch64__)
            DAGBuilder & dagbuilder = DAGBuilder::Instance();
            dagbuilder.finish_thread();
        #endif
        }
    }
    void InitDmaChannel(::google::protobuf::RpcController *controller, const ::rap::DmaUnitInit *request,
                    ::rap::DmaUnitResponse *response, ::google::protobuf::Closure *done) override {
        (void) controller;
        brpc::ClosureGuard done_guard(done);
        //TODO: 相当于可以在这里执行client侧main函数的内容
        DAGBuilder & dagbuilder = DAGBuilder::Instance();
        if (this->if_pipeline) { // pipeline cross device
            dagbuilder.pipeline_init(this->filename, this->feature_num, this->batch_num, this->ele_num, this->if_pipeline);
            BrpcUnit & brpcunit = BrpcUnit::GetInstance();
            brpcunit.export_info(request, response);

            this->backgroundThread = std::thread([&]() {
                LOG(INFO) << "Start pipeline";

                dagbuilder.pipelineunit->start_pipeline();

                std::this_thread::sleep_for(std::chrono::seconds(25));

                LOG(INFO) << "Start end pipeline";

                dagbuilder.pipelineunit->end_pipeline();
            });

            return;
        }

        /* 拓扑排序的跨设备并行 */

        dagbuilder.ProcessUnitInit();
        dagbuilder.set_feature_num(feature_num);
        dagbuilder.set_batch_num(batch_num);
        dagbuilder.set_ele_num(ele_num);
    #if defined(__aarch64__)
        dagbuilder.set_dispatch_type(device_type_t::DEVICE_CPU);
    #else
        dagbuilder.set_dispatch_type(device_type_t::DEVICE_ARM);
    #endif

        // 建图
        dagbuilder.build_from_file(filename);
        
        // 根据op_node的dag图，更新memory node信息（占用大小，ele-num等）
        // 获得memory_offsets 信息
        dagbuilder.memory_assign();

        // 初始化brpc通信单元（其中包含dmauint)
        // 用doca_dma_unit去分配内存(posix)，拿到内存指针
        dagbuilder.memory_region_init();

    #if defined(__aarch64__)
        dagbuilder.dpaProcessUnitInit();
        dagbuilder.task_assign_simple(10, 5, 4);
    #endif

        dagbuilder.server_init(request, response);

    #if defined(__aarch64__)
        dagbuilder.start_thread();
    #endif

        return;
    }

    void TaskDispatch(::google::protobuf::RpcController *controller, const ::rap::TaskSend *request,
                    ::rap::TaskFinish *response, ::google::protobuf::Closure *done) override {
        (void) controller;
        brpc::ClosureGuard done_guard(done);
        DAGBuilder & dagbuilder = DAGBuilder::Instance();

    #if defined(__x86_64__)
        std::vector<int32_t> op_node_list(request->operation_node_list().begin(), 
                                  request->operation_node_list().end());
        dagbuilder.server_task_execute(op_node_list);
    #else // aarch64
        int layer_id = request->layer_id();
        dagbuilder.server_task_execute(layer_id);
    #endif
        response->set_finishflag(1);

        return;
    }

public:
    std::string filename;
    int feature_num;
    int batch_num;
    int ele_num;
    int if_pipeline;
    std::thread backgroundThread;
};
}